# -*- coding: UTF-8 -*-
import tensorflow as tf
import numpy as np
from PIL import Image
import os
import random
import time
import datetime
import utils
LOG_DIR = 'log/cnn1-run-%s' % datetime.datetime.now().strftime('%Y%m%d_%H%M%S')

#训练集，用于训练的验证码图片的文件名
TRAINING_IMAGE_NAME = []
#验证集，用于模型验证的验证码图片的文件名
VALIDATION_IMAGE_NAME = []

X = None
Y = None
x_input = None
keep_prob = None
CNN_PARAM = {}
TRAIN_PARAM = {}
def Init():
    global X,Y,x_input,keep_prob,CNN_PARAM,TRAIN_PARAM
    X = tf.placeholder(tf.float32, [None, TRAIN_PARAM['captcha_image_w'] * TRAIN_PARAM['captcha_image_h']], name='data-input')
    Y = tf.placeholder(tf.float32, [None, TRAIN_PARAM['captcha_len'] * TRAIN_PARAM['char_set_len']], name='label-input')
    x_input = tf.reshape(X, [-1, TRAIN_PARAM['captcha_image_h'], TRAIN_PARAM['captcha_image_w'], 1], name='x-input')

    keep_prob = tf.placeholder(tf.float32)  ##节点保留率

def set_param(param):
    global CNN_PARAM,TRAIN_PARAM
    CNN_PARAM = param
    TRAIN_PARAM = param

def get_next_batch(batchSize=32, trainOrTest='train', step=0):
    global X,Y,x_input,keep_prob,CNN_PARAM,TRAIN_PARAM
    batch_data = np.zeros([batchSize, TRAIN_PARAM['captcha_image_w']*TRAIN_PARAM['captcha_image_h']])
    batch_label = np.zeros([batchSize, TRAIN_PARAM['captcha_len'] * TRAIN_PARAM['char_set_len']])
    fileNameList = TRAINING_IMAGE_NAME
    if trainOrTest == 'validate':
        fileNameList = VALIDATION_IMAGE_NAME

    totalNumber = len(fileNameList)
    indexStart = step*batchSize
    for i in range(batchSize):
        index = (i + indexStart) % totalNumber
        name = fileNameList[index]
        img_data, img_label = utils.get_data_and_label(name,TRAIN_PARAM['captcha_path'],len(name)-4,TRAIN_PARAM['char_set_len'],TRAIN_PARAM['captcha_len'])
        batch_data[i, : ] = img_data
        batch_label[i, : ] = img_label

    return batch_data, batch_label


####################################################################
# 占位符，X和Y分别是输入训练数据和其标签，标签转换成8*10的向量
#X = tf.placeholder(tf.float32, [None, CAPTCHA_IMAGE_WIDHT * CAPTCHA_IMAGE_HEIGHT])
#Y = tf.placeholder(tf.float32, [None, CAPTCHA_LEN * CHAR_SET_LEN])
# 声明dropout占位符变量
#keep_prob = tf.placeholder(tf.float32)  # dropout

def variable_summaries(var):
    global X,Y,x_input,keep_prob,CNN_PARAM,TRAIN_PARAM
    """Attach a lot of summaries to a Tensor (for TensorBoard visualization)."""
    with tf.name_scope('summaries'):
        mean = tf.reduce_mean(var)
        tf.summary.scalar('mean', mean)
        # with tf.name_scope('stddev'):
        #    stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean)))
        # tf.summary.scalar('stddev', stddev)
        # tf.summary.scalar('max', tf.reduce_max(var))
        # tf.summary.scalar('min', tf.reduce_min(var))
        # tf.summary.histogram('histogram', var)

##卷积层 附relu  max_pool drop操作
def conn_layer(w_alpha=0.01, b_alpha=0.1, _keep_prob=None, input=None, last_size=None, cur_size=None):
    global X,Y,x_input,keep_prob,CNN_PARAM,TRAIN_PARAM
    w_c1 = tf.Variable(w_alpha * tf.random_normal([CNN_PARAM['layer_width'], CNN_PARAM['layer_height'], last_size, cur_size]))
    b_c1 = tf.Variable(b_alpha * tf.random_normal([cur_size]))
    conv1 = tf.nn.relu(tf.nn.bias_add(tf.nn.conv2d(input, w_c1, strides=[1, 1, 1, 1], padding='SAME'), b_c1))
    conv1 = tf.nn.max_pool(conv1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
    conv1 = tf.nn.dropout(conv1, keep_prob=_keep_prob)
    return conv1


##对卷积层到全链接层的数据进行变换
def _get_conn_last_size(input):
    global X,Y,x_input,keep_prob
    shape = input.get_shape().as_list()
    dim = 1
    for d in shape[1:]:
        dim *= d
    input = tf.reshape(input, [-1, dim])
    return input, dim


##全链接层
def _fc_layer(w_alpha=0.01, b_alpha=0.1, input=None, last_size=None, cur_size=None):
    global X,Y,x_input,keep_prob,CNN_PARAM,TRAIN_PARAM
    w_d = tf.Variable(w_alpha * tf.random_normal([last_size, cur_size]))
    b_d = tf.Variable(b_alpha * tf.random_normal([cur_size]))
    fc = tf.nn.bias_add(tf.matmul(input, w_d), b_d)
    return fc


##构建前向传播网络
def crack_captcha_cnn(w_alpha=0.01, b_alpha=0.1):
    global X,Y,x_input,keep_prob,CNN_PARAM,TRAIN_PARAM
    x = tf.reshape(X, shape=[-1, TRAIN_PARAM['captcha_image_h'], TRAIN_PARAM['captcha_image_w'], 1])

    conv1 = conn_layer(input=x,_keep_prob=CNN_PARAM['keep_prob'],last_size=1, cur_size=32)
    conv2 = conn_layer(input=conv1,_keep_prob=CNN_PARAM['keep_prob'],last_size=32, cur_size=64)
    conn3 = conn_layer(input=conv2,_keep_prob=CNN_PARAM['keep_prob'],last_size=64, cur_size=64)
    conn4 = conn_layer(input=conn3,_keep_prob=CNN_PARAM['keep_prob'],last_size=64, cur_size=64)
    input, dim = _get_conn_last_size(conn4)
    fc_layer1 = _fc_layer(input=input, last_size=dim, cur_size=1024)
    fc_layer1 = tf.nn.relu(fc_layer1)
    fc_layer1 = tf.nn.dropout(fc_layer1, keep_prob)

    fc_out = _fc_layer(input=fc_layer1, last_size=1024, cur_size=TRAIN_PARAM['captcha_len'] * TRAIN_PARAM['char_set_len'])
    return fc_out

##反向传播
def back_propagation():
    global X,Y,x_input,keep_prob,CNN_PARAM,TRAIN_PARAM
    output = crack_captcha_cnn()
    ##学习率
    loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=Y, logits=output))
    optm = tf.train.AdamOptimizer(learning_rate=0.001).minimize(loss)
    max_idx_p = tf.arg_max(tf.reshape(output, [-1, TRAIN_PARAM['captcha_len'], TRAIN_PARAM['char_set_len']]), 2)
    max_idx_l = tf.arg_max(tf.reshape(Y, [-1, TRAIN_PARAM['captcha_len'], TRAIN_PARAM['char_set_len']]), 2)
    correct_prediction = tf.equal(tf.argmax(max_idx_l, 1), tf.argmax(max_idx_p, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
    #accuracy = tf.metrics.accuracy(labels=max_idx_l,predictions=max_idx_p)
    #accuracy = tf.reduce_mean(tf.cast(tf.equal(max_idx_p, max_idx_l), tf.float32))
    return loss, optm, accuracy


#构建卷积神经网络并训练
def train_first(train_file_list,validate_file_list):
    global X,Y,x_input,keep_prob
    global TRAINING_IMAGE_NAME
    global VALIDATION_IMAGE_NAME
    TRAINING_IMAGE_NAME = train_file_list
    VALIDATION_IMAGE_NAME = validate_file_list
    loss, optm, accuracy = back_propagation()

    with tf.name_scope('loss'):
        variable_summaries(loss)

    # evaluate accuracy
    with tf.name_scope('evaluate_accuracy'):
        variable_summaries(accuracy)

    saver = tf.train.Saver()
    with tf.Session() as sess:
        #tf.summary.scalar('accuracy',acc)
        merged = tf.summary.merge_all()
        train_writer = tf.summary.FileWriter(LOG_DIR + '/train', sess.graph)
        test_writer = tf.summary.FileWriter(LOG_DIR + '/test', sess.graph)
        sess.run(tf.global_variables_initializer())
        steps = 0
        for epoch in range(CNN_PARAM['range_value']):
            train_data, train_label = get_next_batch(CNN_PARAM['batchSize'], 'train', steps)
            step_summary, optm_= sess.run([merged,optm], feed_dict={X : train_data, Y : train_label, keep_prob:0.7})
            #print(step_summary)
            train_writer.add_summary(step_summary, steps)
            if steps % CNN_PARAM['step_value'] == 0:
                test_data, test_label = get_next_batch(CNN_PARAM['batchSize'], 'validate', steps)
                test_summary,acc = sess.run([merged, accuracy], feed_dict={X : test_data, Y : test_label, keep_prob:1.0})
                test_writer.add_summary(test_summary, steps)

                print("steps=%d, accuracy=%f" % (steps, acc))
                if acc > CNN_PARAM['acc']:
                    saver.save(sess, TRAIN_PARAM['model_path']+"crack_captcha.model", global_step=steps)
                    break
            steps += 1

        train_writer.close()
        test_writer.close()

def train_continue(step,train_file_list,validate_file_list):
    global TRAINING_IMAGE_NAME
    global VALIDATION_IMAGE_NAME
    TRAINING_IMAGE_NAME = train_file_list
    VALIDATION_IMAGE_NAME = validate_file_list
    loss, optm, accuracy = back_propagation()

    saver = tf.train.Saver()
    with tf.Session() as sess:
        path = "models/crack_captcha.model-" + str(step)
        saver.restore(sess, path)
        ##36300 36300 0.9325 0.0147698
        while 1:
            batch_x, batch_y = get_next_batch(CNN_PARAM['batchSize'],'train', step)
            _, loss_ = sess.run([optm, loss], feed_dict={X: batch_x, Y: batch_y, keep_prob: 0.75})
            if step % 50 == 0:
                batch_x_test, batch_y_test = get_next_batch(CNN_PARAM['batchSize'], 'validate', step)
                acc = sess.run(accuracy, feed_dict={X: batch_x_test, Y: batch_y_test, keep_prob: 1.})
                print(step, acc, loss_)
                if acc >= 0.99:
                    saver.save(sess, 'models/crack_captcha.model', global_step=step)
                    break
            step += 1

def parse_captcha(file_name,models_file):
    global X,Y,x_input,keep_prob
    image_data = utils.get_test_image_data(file_name)
    output = crack_captcha_cnn()
    result_string = ''
    saver = tf.train.Saver()
    with tf.Session() as sess:
        path = models_file
        saver.restore(sess, path)
        predict = tf.argmax(tf.reshape(output, [-1, TRAIN_PARAM['captcha_len'], TRAIN_PARAM['char_set_len']]), 2)
        text_list = sess.run(predict, feed_dict={X: [image_data], keep_prob: 1})
        result_string = utils.vec2string(text_list[0].tolist())
    saver.close()
    return result_string
